/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.bamboo.mcp.client.service;


import com.bamboo.mcp.client.entity.ToolCallResp;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.ToolResponseMessage;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.model.tool.ToolCallingChatOptions;
import org.springframework.ai.model.tool.ToolCallingManager;
import org.springframework.ai.model.tool.ToolExecutionResult;
import org.springframework.ai.tool.ToolCallbackProvider;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.stereotype.Service;

import java.io.IOException;
import java.time.LocalDateTime;
import java.util.List;
import java.util.Objects;

/**
 * @author brianxiadong
 * Spring AI Alibaba MCP (Model, Chat, Prompt) Service
 */
@Service
public class McpService {

    private final ChatClient chatClient;

    private final ObjectMapper objectMapper;

    private final ToolCallbackProvider tools;

    private final ToolCallingManager toolCallingManager;


    /*
     访问 Ollama 服务模型
    */

    private static final Logger logger = LoggerFactory.getLogger(McpService.class);

    public McpService(
            ObjectMapper objectMapper,
            ToolCallbackProvider tools,
            ToolCallingManager toolCallingManager,
            @Qualifier("openAiChatModel") ChatModel chatModel
            //@Qualifier("ollamaChatModel") ChatModel chatModel
    ) throws IOException {

        this.objectMapper = objectMapper;

        // Initialize chat client with non-blocking configuration
        this.chatClient = ChatClient.builder(chatModel)
                .defaultAdvisors(new SimpleLoggerAdvisor())
                .defaultToolCallbacks(tools)
                .build();
        this.tools = tools;
        this.toolCallingManager = toolCallingManager;


    }

    public String chat(String prompt) {
        String llmCallResponse = "";
        // manual run tools flag
        ChatOptions chatOptions = ToolCallingChatOptions.builder()
                .toolCallbacks(tools.getToolCallbacks())
                .build();

        ChatResponse response = chatClient.prompt(new Prompt(prompt, chatOptions)).call().chatResponse();

        logger.debug("ChatResponse: {}", response.toString());
        assert response != null;
        List<AssistantMessage.ToolCall> toolCalls = response.getResult().getOutput().getToolCalls();
        logger.debug("ToolCalls: {}", toolCalls);
        llmCallResponse = response.getResult().getOutput().getText();
        logger.debug("Response by LLM: {}", llmCallResponse);

        // execute tools with no chat memory messages.
        ToolCallResp tcr = ToolCallResp.TCR();
        if (!toolCalls.isEmpty()) {

            tcr = ToolCallResp.startExecute(
                    llmCallResponse,
                    toolCalls.get(0).name(),
                    toolCalls.get(0).arguments()
            );
            tcr.setToolParameters(toolCalls.get(0).arguments());
            logger.debug("Start ToolCallResp: {}", tcr);
            ToolExecutionResult toolExecutionResult = null;

            try {
                toolExecutionResult = toolCallingManager.executeToolCalls(new Prompt(prompt, chatOptions), response);

                tcr.setToolEndTime(LocalDateTime.now());
            } catch (Exception e) {

                tcr.setStatus(ToolCallResp.ToolState.FAILURE);
                tcr.setErrorMessage(e.getMessage());
                tcr.setToolEndTime(LocalDateTime.now());
                tcr.setToolCostTime((long) (tcr.getToolEndTime().getNano() - tcr.getToolStartTime().getNano()));
                logger.error("Error ToolCallResp: {}, msg: {}", tcr, e.getMessage());
                // throw new RuntimeException("Tool execution failed, please check the logs for details.");
            }

            if (Objects.nonNull(toolExecutionResult)) {
//				ToolResponseMessage toolResponseMessage = (ToolResponseMessage) toolExecutionResult.conversationHistory()
//						.get(toolExecutionResult.conversationHistory().size() - 1);
//				llmCallResponse = toolResponseMessage.getResponses().get(0).responseData();
                ChatResponse finalResponse = chatClient.prompt().messages(toolExecutionResult.conversationHistory()).call().chatResponse();
                if (finalResponse != null) {
                    llmCallResponse = finalResponse.getResult().getOutput().getText();
                }

                StringBuilder sb = new StringBuilder();
                toolExecutionResult.conversationHistory().stream()
                        .filter(message -> message instanceof ToolResponseMessage)
                        .forEach(message -> {
                            ToolResponseMessage toolResponseMessage = (ToolResponseMessage) message;
                            toolResponseMessage.getResponses().forEach(tooResponse -> {
                                sb.append(tooResponse.responseData());
                            });
                        });
                tcr.setToolResponse(sb.toString());
            }

            tcr.setStatus(ToolCallResp.ToolState.SUCCESS);
            tcr.setToolResult(llmCallResponse);
            tcr.setToolCostTime((long) (tcr.getToolEndTime().getNano() - tcr.getToolStartTime().getNano()));
            logger.debug("End ToolCallResp: {}", tcr);
            return llmCallResponse;
        } else {
            logger.debug("ToolCalls is empty, no tool execution needed.");
            tcr.setToolResult(llmCallResponse);
        }
        logger.debug("End ToolCallResp: {}", tcr);
        return llmCallResponse;
    }

}

